#include <memory>
#include "tensorflow/core/framework/fake_input.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/node_def_builder.h"
#include "tensorflow/core/framework/shape_inference.h"
#include "tf_adapter/kernels/aicpu/npu_embedding_ops.cc"
#include "gtest/gtest.h"

namespace tensorflow {
namespace {
PartialTensorShape TShape(std::initializer_list<int64> dims) {
  return PartialTensorShape(dims);
}

FakeInputFunctor FakeInputStub(DataType dt) {
  return [dt](const OpDef &op_def, int in_index, const NodeDef &node_def,
              NodeDefBuilder *builder) {
    char c = 'a' + (in_index % 26);
    string in_node = string(&c, 1);
    builder->Input(in_node, 0, dt);
    return Status::OK();
  };
}

FakeInputFunctor FakeInputStubList(DataType dt) {
  return [dt](const OpDef &op_def, int in_index, const NodeDef &node_def,
              NodeDefBuilder *builder) {
    char c = 'a' + (in_index % 26);
    string in_node = string(&c, 1);
    builder->Input({{in_node, {}, dt}});
    return Status::OK();
  };
}

class NpuCpuOpTest : public testing::Test {
 protected:
  virtual void SetUp() {}
  virtual void TearDown() {}
};

TEST(EmbeddingOpsTest, TestEmbeddingTableFind02) {
    const OpRegistrationData *reg;
    TF_CHECK_OK(OpRegistry::Global()->LookUp("EmbeddingTableFind", &reg));
    OpDef op_def = reg->op_def;
    NodeDef def;
    TF_CHECK_OK(NodeDefBuilder("dummy", &op_def)
                    .Attr("embedding_dim", {4})
                    .Input(FakeInputStub(DT_INT32))
                    .Input(FakeInputStub(DT_INT64))
                    .Finalize(&def));

    shape_inference::InferenceContext c(
        0, &def, op_def,
        {TShape({1}), TShape({16})},
        {}, {}, {});
    ASSERT_TRUE(reg->shape_inference_fn(&c).ok());
}

TEST(EmbeddingOpsTest, TestEmbeddingFeatureMappingShapeInfer) {
  const OpRegistrationData* reg;
  TF_CHECK_OK(OpRegistry::Global()->LookUp("EmbeddingFeatureMapping", &reg));
  OpDef op_def = reg->op_def;
  NodeDef def;
  TF_CHECK_OK(NodeDefBuilder("dummy", &op_def)
                  .Input(FakeInputStub(DT_INT64))
                  .Finalize(&def));
  shape_inference::InferenceContext c(0, &def, op_def, {TShape({2, 2, 3, 4})}, {}, {}, {});
  ASSERT_TRUE(reg->shape_inference_fn(&c).ok());
}

TEST(EmbeddingOpsTest, TestEmbeddingFeatureMappingV2ShapeInfer) {
  const OpRegistrationData* reg;
  TF_CHECK_OK(OpRegistry::Global()->LookUp("EmbeddingFeatureMappingV2", &reg));
  OpDef op_def = reg->op_def;
  NodeDef def;
  TF_CHECK_OK(NodeDefBuilder("dummy", &op_def)
                    .Attr("table_total_size", {1})
                    .Attr("table_actual_size", {1})
                    .Input(FakeInputStub(DT_STRING))
                    .Input(FakeInputStub(DT_INT64))
                    .Finalize(&def));
  shape_inference::InferenceContext c(
        0, &def, op_def,
        {TShape({}), TShape({6})},
        {}, {}, {});
  ASSERT_TRUE(reg->shape_inference_fn(&c).ok());
}

TEST(EmbeddingOpsTest, TestEmbeddingFeatureMappingFindShapeInfer) {
    const OpRegistrationData *reg;
    TF_CHECK_OK(OpRegistry::Global()->LookUp("EmbeddingFeatureMappingFind", &reg));
    OpDef op_def = reg->op_def;
    NodeDef def;
    TF_CHECK_OK(NodeDefBuilder("dummy", &op_def)
                    .Attr("num", 1)
                    .Input(FakeInputStub(DT_STRING))
                    .Input(FakeInputStub(DT_INT64))
                    .Finalize(&def));
    shape_inference::InferenceContext c(
        0, &def, op_def,
        {TShape({}), TShape({})},
        {}, {}, {});
    ASSERT_TRUE(reg->shape_inference_fn(&c).ok());
}

TEST(EmbeddingOpsTest, TestEmbeddingFeatureMappingInsertShapeInfer) {
    const OpRegistrationData *reg;
    TF_CHECK_OK(OpRegistry::Global()->LookUp("EmbeddingFeatureMappingInsert", &reg));
    OpDef op_def = reg->op_def;
    NodeDef def;
    TF_CHECK_OK(NodeDefBuilder("dummy", &op_def)
                    .Input(FakeInputStub(DT_STRING))
                    .Input(FakeInputStubList(DT_INT64))
                    .Input(FakeInputStubList(DT_INT32))
                    .Finalize(&def));
    shape_inference::InferenceContext c(
        0, &def, op_def,
        {TShape({}), TShape({6}), TShape({6})},
        {}, {}, {});
    ASSERT_TRUE(reg->shape_inference_fn(&c).ok());
}

TEST(EmbeddingOpsTest, TestEmbeddingFeatureMappingExportShapeInfer) {
    const OpRegistrationData *reg;
    TF_CHECK_OK(OpRegistry::Global()->LookUp("EmbeddingFeatureMappingExport", &reg));
    OpDef op_def = reg->op_def;
    NodeDef def;
    TF_CHECK_OK(NodeDefBuilder("dummy", &op_def)
                    .Input(FakeInputStub(DT_STRING))
                    .Input(FakeInputStub(DT_STRING))
                    .Input(FakeInputStub(DT_INT64))
                    .Input(FakeInputStub(DT_FLOAT))
                    .Input(FakeInputStubList(DT_INT64))
                    .Input(FakeInputStubList(DT_INT32))
                    .Finalize(&def));
    shape_inference::InferenceContext c(
        0, &def, op_def,
        {TShape({}), TShape({}), TShape({6}), TShape({6}), TShape({6})},
        {}, {}, {});
    ASSERT_TRUE(reg->shape_inference_fn(&c).ok());
}

TEST(EmbeddingOpsTest, TestEmbeddingFeatureMappingImportShapeInfer) {
    const OpRegistrationData *reg;
    TF_CHECK_OK(OpRegistry::Global()->LookUp("EmbeddingFeatureMappingImport", &reg));
    OpDef op_def = reg->op_def;
    NodeDef def;
    TF_CHECK_OK(NodeDefBuilder("dummy", &op_def)
                    .Attr("embedding_dim", 4)
                    .Attr("only_offset_flag", 1)
                    .Attr("num", 1)
                    .Input(FakeInputStub(DT_STRING))
                    .Input(FakeInputStub(DT_STRING))
                    .Input(FakeInputStub(DT_INT64))
                    .Input(FakeInputStub(DT_INT64))
                    .Finalize(&def));
    shape_inference::InferenceContext c(
        0, &def, op_def,
        {TShape({}), TShape({}), TShape({6})},
        {}, {}, {});
    ASSERT_TRUE(reg->shape_inference_fn(&c).ok());
}

TEST(EmbeddingOpsTest, TestEmbeddingFeatureMappingTableSizeShapeInfer) {
    const OpRegistrationData *reg;
    TF_CHECK_OK(OpRegistry::Global()->LookUp("EmbeddingFeatureMappingTableSize", &reg));
    OpDef op_def = reg->op_def;
    NodeDef def;
    TF_CHECK_OK(NodeDefBuilder("dummy", &op_def)
                    .Input(FakeInputStub(DT_STRING))
                    .Finalize(&def));
    shape_inference::InferenceContext c(
        0, &def, op_def,
        {TShape({})},
        {}, {}, {});
    ASSERT_TRUE(reg->shape_inference_fn(&c).ok());
}

TEST(EmbeddingOpsTest, TestEmbeddingFeatureMappingFileSizeShapeInfer) {
    const OpRegistrationData *reg;
    TF_CHECK_OK(OpRegistry::Global()->LookUp("EmbeddingFeatureMappingFileSize", &reg));
    OpDef op_def = reg->op_def;
    NodeDef def;
    TF_CHECK_OK(NodeDefBuilder("dummy", &op_def)
                    .Attr("embedding_dim", 4)
                    .Attr("only_offset_flag", 1)
                    .Input(FakeInputStub(DT_STRING))
                    .Input(FakeInputStub(DT_STRING))
                    .Input(FakeInputStub(DT_INT64))
                    .Finalize(&def));
    shape_inference::InferenceContext c(
        0, &def, op_def,
        {TShape({}), TShape({})},
        {}, {}, {});
    ASSERT_TRUE(reg->shape_inference_fn(&c).ok());
}

TEST(EmbeddingOpsTest, InitEmbeddingHashmapV2ShapeInfer) {
  const OpRegistrationData* reg;
  TF_CHECK_OK(OpRegistry::Global()->LookUp("InitEmbeddingHashmapV2", &reg));
  OpDef op_def = reg->op_def;
  NodeDef def;
  TF_CHECK_OK(NodeDefBuilder("dummy", &op_def)
                    .Attr("bucket_size", 10)
                    .Attr("load_factor", 80)
                    .Attr("embedding_dim", 2)
                    .Attr("dtype", DT_FLOAT)
                    .Input(FakeInputStub(DT_INT32))
                    .Finalize(&def));
  shape_inference::InferenceContext c(
        0, &def, op_def,
        {TShape({6})},
        {}, {}, {});
  ASSERT_TRUE(reg->shape_inference_fn(&c).ok());
}

TEST(EmbeddingOpsTest, DeinitEmbeddingHashmapV2ShapeInfer) {
  const OpRegistrationData* reg;
  TF_CHECK_OK(OpRegistry::Global()->LookUp("DeinitEmbeddingHashmapV2", &reg));
  OpDef op_def = reg->op_def;
  NodeDef def;
  TF_CHECK_OK(NodeDefBuilder("dummy", &op_def)
                    .Input(FakeInputStub(DT_INT32))
                    .Finalize(&def));
  shape_inference::InferenceContext c(
        0, &def, op_def,
        {TShape({6})},
        {}, {}, {});
  ASSERT_TRUE(reg->shape_inference_fn(&c).ok());
}

TEST(EmbeddingOpsTest, TableToResourceV2ShapeInfer) {
  const OpRegistrationData* reg;
  TF_CHECK_OK(OpRegistry::Global()->LookUp("TableToResourceV2", &reg));
  OpDef op_def = reg->op_def;
  NodeDef def;
  TF_CHECK_OK(NodeDefBuilder("dummy", &op_def)
                    .Input(FakeInputStub(DT_INT32))
                    .Finalize(&def));
  shape_inference::InferenceContext c(
        0, &def, op_def,
        {TShape({6})},
        {}, {}, {});
  ASSERT_TRUE(reg->shape_inference_fn(&c).ok());
}

TEST(EmbeddingOpsTest, EmbeddingHashmapImportShapeInfer) {
    const OpRegistrationData *reg;
    TF_CHECK_OK(OpRegistry::Global()->LookUp("EmbeddingHashmapImport", &reg));
    OpDef op_def = reg->op_def;
    NodeDef def;
    TF_CHECK_OK(NodeDefBuilder("dummy", &op_def)
                    .Attr("embedding_dim", 4)
                    .Attr("num", 1)
                    .Input(FakeInputStub(DT_STRING))
                    .Input(FakeInputStub(DT_INT32))
                    .Input(FakeInputStub(DT_INT64))
                    .Input(FakeInputStub(DT_STRING))
                    .Input(FakeInputStub(DT_INT64))
                    .Finalize(&def));
    shape_inference::InferenceContext c(
        0, &def, op_def,
        {TShape({}), TShape({}), TShape({6})},
        {}, {}, {});
    ASSERT_TRUE(reg->shape_inference_fn(&c).ok());
}
}
}
